package com.mjf;

import com.mjf.util.DataLoader;
import ml.dmlc.xgboost4j.java.Booster;
import ml.dmlc.xgboost4j.java.DMatrix;
import ml.dmlc.xgboost4j.java.XGBoost;
import ml.dmlc.xgboost4j.java.XGBoostError;

import java.io.File;
import java.io.IOException;
import java.io.PrintWriter;
import java.util.Arrays;
import java.util.HashMap;

public class XGBoostJavaExample {
    public static void main(String[] args) throws XGBoostError, IOException {

        // 训练数据集
        DMatrix trainData = new DMatrix("./input/agaricus.txt.train");

        // 测试数据集
        DMatrix testData = new DMatrix("./input/agaricus.txt.test");

        // 模型参数
        HashMap<String, Object> params = new HashMap<String, Object>();
        params.put("eta", 1.0);
        params.put("max_depth", 2);
        params.put("silent", 1);
        params.put("objective", "binary:logistic");

        // 设置训练集与测试集
        HashMap<String, DMatrix> watches = new HashMap<String, DMatrix>();
        watches.put("train", trainData);
        watches.put("test", testData);

        // 迭代次数
        int round = 2;

        // 训练模型
        Booster booster = XGBoost.train(trainData, params, round, watches, null, null);

        // 预测
        float[][] predicts = booster.predict(testData);

        // 模型保存路径
        File file = new File("./model");
        if (!file.exists()) {
            file.mkdirs();
        }

        // 将模型保存到本地
        booster.saveModel("./model/xgb.model");

        // 获取带有特征图的转储模型
        String[] modelInfos = booster.getModelDump("./input/featmap.txt", false);
        saveDumpModel("./model/dump.raw.txt", modelInfos);

        // 将测试数据以二进制缓存
        testData.saveBinary("./model/dtest.buffer");

        // 重新加载模型和数据
        Booster booster2 = XGBoost.loadModel("./model/xgb.model");
        DMatrix testData2 = new DMatrix("./model/dtest.buffer");

        // 第二次预测
        float[][] predicts2 = booster2.predict(testData2);

        // 检查两次预测结果是否相等
        System.out.println(checkPredicts(predicts, predicts2));

        System.out.println("从CSR稀疏数据开始构建dmatrix...");
        
        // 从CSR稀疏数据开始构建dmatrix
        DataLoader.CSRSparseData spData = DataLoader.loadSVMFile("./input/agaricus.txt.train");

        DMatrix trainData2 = new DMatrix(spData.rowHeaders, spData.colIndex, spData.data,
                DMatrix.SparseType.CSR, 127);
        trainData2.setLabel(spData.labels);

        // 指定训练集与测试集
        HashMap<String, DMatrix> watches2 = new HashMap<String, DMatrix>();
        watches2.put("train", trainData2);
        watches2.put("test", testData2);

        Booster booster3 = XGBoost.train(trainData2, params, round, watches2, null, null);
        float[][] predicts3 = booster3.predict(testData2);

        // 检查两次预测结果是否相等
        System.out.println(checkPredicts(predicts, predicts3));


    }

    public static boolean checkPredicts(float[][] fPredicts, float[][] sPredicts) {
        if (fPredicts.length != sPredicts.length) {
            return false;
        }

        for (int i = 0; i < fPredicts.length; i++) {
            if (!Arrays.equals(fPredicts[i], sPredicts[i])) {
                return false;
            }
        }

        return true;
    }

    public static void saveDumpModel(String modelPath, String[] modelInfos) throws IOException {
        try {
            PrintWriter writer = new PrintWriter(modelPath, "UTF-8");
            for (int i = 0; i < modelInfos.length; ++i) {
                writer.print("booster[" + i + "]:\n");
                writer.print(modelInfos[i]);
            }
            writer.close();
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

}
